#include "dce.h"

#include <iostream>
#include <ostream>
#include <string>

void DeadCodeEliminator::BuildCD(IRFunction* func) {
    CD.clear(), RCD.clear();
    T.Build(func->exit, true);
    for (auto u : T.nodes) {
        for (auto v : T.DF[u]) {
            CD[v].push_back(u);
            RCD[u].push_back(v);
        }
    }
}

void DeadCodeEliminator::RunFunc(IRFunction* func) {
    BuildCD(func);

    std::unordered_set<std::string> argMem;
    for (auto i : func->argMemEffects)
        argMem.insert(func->args[i]->memBase->name);

    std::unordered_map<std::string, std::vector<IR::IRBase*>> changer;
    std::unordered_map<IR::IRBase*, std::vector<IR::IRBase*>> definer;
    std::unordered_set<IR::IRBase*> live;
    live.insert(func->exit->GetTerminator());
    for (auto u : T.nodes) {
        for (auto ir : u->GetAllIR()) {
            auto eff = ir->Effects();
            if (eff.io) live.insert(ir);
            for (auto [name, var] : eff.var) {
                changer[name].push_back(ir);
                if (var.isGlobal) live.insert(ir);
            }
            for (auto [name, var] : eff.mem) {
                changer[name].push_back(ir);
                if (var.isGlobal) live.insert(ir);
                if (argMem.count(name)) live.insert(ir);
            }
        }
    }
    for (auto u : T.nodes) {
        for (auto ir : u->GetAllIR()) {
            for (auto var : ir->UsedVar()) {
                definer[ir].insert(definer[ir].end(),
                                   changer[var->name].begin(),
                                   changer[var->name].end());
                if (var->type.isPtr) {
                    auto name = var->memBase->name;
                    definer[ir].insert(definer[ir].end(), changer[name].begin(),
                                       changer[name].end());
                }
            }
        }
    }

    std::queue<IR::IRBase*> q;
    for (auto u : live) q.push(u);
    while (!q.empty()) {
        auto u = q.front();
        q.pop();
        for (auto v : definer[u]) {
            if (!live.count(v)) {
                live.insert(v);
                q.push(v);
            }
        }
        for (auto block : RCD[u->parent]) {
            auto v = block->GetTerminator();
            if (!live.count(v)) {
                live.insert(v);
                q.push(v);
            }
        }
    }

    for (auto u : T.nodes) {
        for (auto it = u->phis.begin(); it != u->phis.end();) {
            if (!live.count(*it)) {
                // (*it)->Print(), std::cout << std::endl;
                it = u->phis.erase(it);
            } else
                ++it;
        }
        for (auto it = u->irList.begin(); it != std::prev(u->irList.end());) {
            if (!live.count(*it)) {
                // (*it)->Print(), std::cout << std::endl;
                it = u->irList.erase(it);
            } else
                ++it;
        }
        if (auto br = u->GetTerminator()->As<IR::Branch>()) {
            if (live.count(br)) continue;
            u->RemoveBranch();
        }
    }

    func->Simplify();
    func->SetParent();
    func->Validate();
}

void* DeadCodeEliminator::Run(void* input) {
    auto module = (IRModule*)input;
    for (auto func : module->globalFuncs) {
        RunFunc(func);
    }
    return module;
}
